import base64

import cv2
import numpy as np
from flask import Flask, request, jsonify
from paddlex import create_pipeline

app = Flask(__name__)

# 创建 pipeline
# pipeline = create_pipeline(pipeline="./config/PP-ChatOCRv3-doc.yaml")

pipeline = create_pipeline(
    pipeline="PP-ChatOCRv3-doc",
    llm_name="ernie-4.0",
    llm_params={"api_type": "aistudio", "access_token": "9b0c7f625a8badfbc19691472e336ec582ea04a4"},
    layout_model="./official_models/PicoDet-L_layout_3cls",  # 版面区域检测模块模型s
    table_model="./official_models/SLANet",  # 表格识别模型
    text_det_model="./official_models/PP-OCRv4_server_det",  # 文本检测模块模型
    text_rec_model="./official_models/PP-OCRv4_mobile_rec",  # 文本识别模块模型
    seal_text_det_model="./official_models/PP-OCRv4_mobile_seal_det",  # 印章文本检测
    # doc_image_ori_cls_model="PP-LCNet_x1_0_doc_ori",  # 文档图像方向分类模块模型
    text_det_batch_size=1,  # 该参数只能是1
    text_rec_batch_size=5
)


def base64_to_ndarray(base64_str):
    # 将 base64 字符串解码为二进制图像数据
    img_data = base64.b64decode(base64_str)
    # 转换为 numpy 数组并解码为 OpenCV 格式的图像
    nparr = np.frombuffer(img_data, np.uint8)
    image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    if image is None:
        raise ValueError("图片解码失败，请检查传入的 base64 字符串")
    return image


def get_Kword(visual_info):
    text_block = visual_info['ocr_text'][0]['words in text block']
    kIdCard = [
        "请给出正确的身份证信息,包含:姓名、性别、民族、出生、身份证号、住址，身份证号必须是完整的18位数字,格式为 {type:ID,name:xxx,sex:xxx,nation:xxx ,birth:xxx,idNo:xxxxxxxxxxxxxxxxx ,address:xxx}"]

    kBankCard = [
        "请给出正确的银行卡号和银行名称信息,银行卡号必须是完整的19位数字或者16位数字,"
        "由前面6位数字加后面13位数字构成的19位数字或者由连续的4位构成的16位数字并且连续的4位数字不会重复，"
        "银行名称必须是中文，格式为 {type:bank,cardNo:xxxxxx xxxxxxxxxxxxx or xxxx xxxx xxxx xxxx，cardName:招商银行} carNO 即银行卡号 必须是19位或者16位数字"]
    kLicense = [
        "请给出正确的营业执照信息,包含:营业执照注册号即（统一社会信用代码）、公司名称即（名称）、公司类型、法定代表人、经营范围、注册资本、成立日期、营业期限、住所即（公司地址），营业执照注册号必须是完整的18位数字,"
        "格式为 {type:license,licenseNo:xxxxxxxxxxxxxxxxx,companyName:xxx,companyType:xxxx,corporation:xxx,scope:xxx,registeredCapital:xxx,establishmentDate:xxx,businessPeriod:xxx,companyAddress:xxx}"]

    # kWeightTicket = [
    #     "从图片中提取所有文本信息，"
    #     "表格第一行包括、供应商、车号、日期、皮重过磅时间、皮重过磅重量、皮重过磅员、"
    #     "表格第二行毛重过磅时间、毛重过磅重量、毛重过磅员、净重、"
    #     "表格第三行货物信息 包括 货物名称、结算重量等。"
    #     "以结构化的格式返回数据，"
    #     "格式为 {供应商:xxx,车号:xxx,日期:xxx,皮重过磅时间:xxx,皮重过磅重量:xxx,皮重过磅员:xxx,毛重过磅时间:xxx,毛重过磅重量:xxx,毛重过磅员:xxx,净重:xxx,货物名称:xxx,"
    #     "结算重量:xxx}"
    #     ""]
    kWeightTicket = ["从图片中提取所有文本信息,输出markdown格式的表格数据"]
    ## 包含的关键词

    weight_ticket_key_word = ["过磅单"]

    bank_key_word = ["Debit", "DEBIT", "ATM", "银联", "BOC", "CCB", "ABC", "ICBC", "BOCOM", "CMB", "CIB", "CMBC", "HXB",
                     "PAB",
                     "CEB", "CNCB", "SPDB"]

    id_key_word = ["公民身份号码"]

    license_key_word = ["法定代表人"]

    bank_advice = ["回单"]
    # 判断文本中是否包含关键词

    keyword = []
    # 判断文本中是否包含关键词
    if all(key in text_block for key in license_key_word):
        keyword = kLicense
    elif all(key in text_block for key in id_key_word):
        keyword = kIdCard
    elif all(key in text_block for key in weight_ticket_key_word):
        keyword = kWeightTicket
    elif all(key in text_block for key in bank_advice):
        keyword = ""
    else:
        keyword = kBankCard

    return keyword


@app.route('/api/rec', methods=['POST'])
def predict():
    # 从请求中获取图片的 base64 字符串
    data = request.json.get("image")
    if not data:
        return jsonify({"error": "No image data provided"}), 400

    try:
        # 将 base64 图片字符串转换为 numpy.ndarray
        image = base64_to_ndarray(data)
    except Exception as e:
        return jsonify({"error": str(e)}), 400

    # 使用 pipeline 进行视觉预测
    try:
        visual_result, visual_info = pipeline.visual_predict(image)
    except TypeError as e:
        return jsonify({"error": str(e)}), 500

    vector = pipeline.build_vector(visual_info=visual_info)

    keyword = get_Kword(visual_info)
    print(keyword)
    chat_result = pipeline.chat(
        key_list=keyword,
        visual_info=visual_info,
        vector=vector,
    )
    print(chat_result)
    print("--------------------------------------")
    print(chat_result["chat_res"])

    return chat_result["chat_res"]


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=6500, debug=True)
